#! /usr/bin/env python3

import rospy
import rosbag
from grid_map import *
from utils import *
import os
import os.path as osp
from nav_msgs.msg import OccupancyGrid


BASE_PATH = osp.abspath('.')
BAG_PATH = osp.join(BASE_PATH, 'bags', 'stage_2.bag')

P_prior = 0.5
P_occ = 0.9
P_free = 0.3
RESOLUTION = 0.04
dir_pointer_len = 0.15
map_x_lim = [-4, 4]
map_y_lim = [-4, 4]

if __name__ == "__main__":
    mapping_node = rospy.init_node("mapping", anonymous=False, disable_signals=True)
    pub = rospy.Publisher("map", OccupancyGrid, queue_size=10)
    rate = rospy.Rate(10)
    try:
        bag = rosbag.Bag(BAG_PATH)
        N_odom = 0
        N_scan = 0
        N_paired = 0
        N_unpaired = 0

        # msg_handler for offline mapping
        msgHandler = MessageHandler()

        # grid map
        grid_map = GridMap(map_x_lim, map_y_lim, RESOLUTION, P_prior)

        # TODO: draw the map

        map = OccupancyGrid()
        map.info.resolution = RESOLUTION
        map.info.width = int((map_x_lim[1] - map_x_lim[0]) / RESOLUTION)
        map.info.height = int((map_x_lim[1] - map_x_lim[0]) / RESOLUTION)
        map.info.origin.position.x = 0
        map.info.origin.position.y = 0
        map.info.origin.position.z = 0
        map.info.origin.orientation.x = 0
        map.info.origin.orientation.y = 0
        map.info.origin.orientation.z = 0
        map.info.origin.orientation.w = 1
        map.header.stamp = rospy.Time.now()
        map.header.frame_id = "odom"

        # read messages from a rosbag
        for topic, msg, time_stamp in bag.read_messages():
            if topic == "/scan":
                N_scan += 1
                # print(np.any(np.array(msg.ranges) == np.inf))
            elif topic == "/odom":
                N_odom += 1

            msgHandler.accept_message(topic, msg)
            if msgHandler.pair_check() == False:
                N_unpaired += 1
                continue
            else:
                N_paired += 1

            scan_msg, odom_msg = msgHandler.process_paired_messages()

            # process these messages
            x_odom, y_odom, theta = convert_odom(odom_msg)
            distances, angles, information = convert_lidar_scan_to_local(scan_msg)
            x_in, y_in = robotic_2_inertial(x_odom, y_odom, theta, distances, angles)

            # mapping
            x1, y1 = grid_map.discretize(x_odom, y_odom)
            X2, Y2 = [], []
            for x, y, dist in zip(x_in, y_in, distances):
                x2, y2 = grid_map.discretize(x, y)

                for x_bres, y_bres in bresenham(x1, y1, x2, y2):
                    grid_map.update(x_bres, y_bres, P_free)
                if dist < scan_msg.range_max:
                    grid_map.update(x2, y2, P_occ)

            p = grid_map.retrieve_p()
            data = np.flip(p, axis=0).reshape(-1).astype(np.int32) * 100
            map.data = data.tolist()
            pub.publish(map)


    except KeyboardInterrupt:
        rospy.loginfo("\nMapping terminated manually")
        pass